// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/compute/light_array_internal.h>

#include <type_traits>

#include <nebula/bits/bitmap_ops.h>
#include <nebula/numeric/int_util_overflow.h>
#include <turbo/base/macros.h>

namespace nebula::compute {

    KeyColumnArray::KeyColumnArray(const KeyColumnMetadata &metadata, int64_t length,
                                   const uint8_t *validity_buffer,
                                   const uint8_t *fixed_length_buffer,
                                   const uint8_t *var_length_buffer, int bit_offset_validity,
                                   int bit_offset_fixed) {
        static_assert(
                std::is_trivial_v<KeyColumnArray> && std::is_standard_layout_v<KeyColumnArray>,
                "This class was intended to be a POD type");
        metadata_ = metadata;
        length_ = length;
        buffers_[kValidityBuffer] = validity_buffer;
        buffers_[kFixedLengthBuffer] = fixed_length_buffer;
        buffers_[kVariableLengthBuffer] = var_length_buffer;
        mutable_buffers_[kValidityBuffer] = mutable_buffers_[kFixedLengthBuffer] =
        mutable_buffers_[kVariableLengthBuffer] = nullptr;
        bit_offset_[kValidityBuffer] = bit_offset_validity;
        bit_offset_[kFixedLengthBuffer] = bit_offset_fixed;
    }

    KeyColumnArray::KeyColumnArray(const KeyColumnMetadata &metadata, int64_t length,
                                   uint8_t *validity_buffer, uint8_t *fixed_length_buffer,
                                   uint8_t *var_length_buffer, int bit_offset_validity,
                                   int bit_offset_fixed) {
        metadata_ = metadata;
        length_ = length;
        buffers_[kValidityBuffer] = mutable_buffers_[kValidityBuffer] = validity_buffer;
        buffers_[kFixedLengthBuffer] = mutable_buffers_[kFixedLengthBuffer] =
                fixed_length_buffer;
        buffers_[kVariableLengthBuffer] = mutable_buffers_[kVariableLengthBuffer] =
                var_length_buffer;
        bit_offset_[kValidityBuffer] = bit_offset_validity;
        bit_offset_[kFixedLengthBuffer] = bit_offset_fixed;
    }

    KeyColumnArray KeyColumnArray::WithBufferFrom(const KeyColumnArray &other,
                                                  int buffer_id_to_replace) const {
        KeyColumnArray copy = *this;
        copy.mutable_buffers_[buffer_id_to_replace] =
                other.mutable_buffers_[buffer_id_to_replace];
        copy.buffers_[buffer_id_to_replace] = other.buffers_[buffer_id_to_replace];
        if (buffer_id_to_replace < kMaxBuffers - 1) {
            copy.bit_offset_[buffer_id_to_replace] = other.bit_offset_[buffer_id_to_replace];
        }
        return copy;
    }

    KeyColumnArray KeyColumnArray::with_metadata(const KeyColumnMetadata &metadata) const {
        KeyColumnArray copy = *this;
        copy.metadata_ = metadata;
        return copy;
    }

    KeyColumnArray KeyColumnArray::slice(int64_t offset, int64_t length) const {
        KeyColumnArray sliced;
        sliced.metadata_ = metadata_;
        sliced.length_ = length;
        uint32_t fixed_size = metadata_.fixed_length;

        sliced.buffers_[0] =
                buffers_[0] ? buffers_[0] + (bit_offset_[0] + offset) / 8 : nullptr;
        sliced.mutable_buffers_[0] =
                mutable_buffers_[0] ? mutable_buffers_[0] + (bit_offset_[0] + offset) / 8 : nullptr;
        sliced.bit_offset_[0] = (bit_offset_[0] + offset) % 8;

        if (metadata_.fixed_length == 0 && !metadata_.is_null_type) {
                    DKCHECK(is_bool_type()) << "Expected BOOL type but got a different type.";
            sliced.buffers_[1] =
                    buffers_[1] ? buffers_[1] + (bit_offset_[1] + offset) / 8 : nullptr;
            sliced.mutable_buffers_[1] = mutable_buffers_[1]
                                         ? mutable_buffers_[1] + (bit_offset_[1] + offset) / 8
                                         : nullptr;
            sliced.bit_offset_[1] = (bit_offset_[1] + offset) % 8;
        } else if (metadata_.fixed_length > 0) {
                    DKCHECK(is_binary_type() || is_large_binary_type() || is_fixed_width_types())
                    << "Expected (LARGE) BINARY or FIXED WIDTH type but got a different type.";
            sliced.buffers_[1] = buffers_[1] ? buffers_[1] + offset * fixed_size : nullptr;
            sliced.mutable_buffers_[1] =
                    mutable_buffers_[1] ? mutable_buffers_[1] + offset * fixed_size : nullptr;
            sliced.bit_offset_[1] = 0;
        } else {
                    DKCHECK(is_null_type()) << "Expected Null type but got a different type.";
        }

        sliced.buffers_[2] = buffers_[2];
        sliced.mutable_buffers_[2] = mutable_buffers_[2];
        return sliced;
    }

    turbo::Result<KeyColumnMetadata> ColumnMetadataFromDataType(
            const std::shared_ptr<DataType> &type) {
        const bool is_extension = type->id() == Type::EXTENSION;
        const std::shared_ptr<DataType> &typ =
                is_extension
                ? turbo::checked_pointer_cast<ExtensionType>(type->get_shared_ptr())
                        ->storage_type()
                : type;

        if (typ->id() == Type::DICTIONARY) {
            auto bit_width =
                    turbo::checked_cast<const FixedWidthType &>(*typ).bit_width();
                    DKCHECK(bit_width % 8 == 0);
            return KeyColumnMetadata(true, bit_width / 8);
        }
        if (typ->id() == Type::BOOL) {
            return KeyColumnMetadata(true, 0);
        }
        if (is_fixed_width(typ->id())) {
            return KeyColumnMetadata(
                    true, turbo::checked_cast<const FixedWidthType &>(*typ).bit_width() / 8);
        }
        if (is_binary_like(typ->id())) {
            return KeyColumnMetadata(false, sizeof(uint32_t));
        }
        if (is_large_binary_like(typ->id())) {
            return KeyColumnMetadata(false, sizeof(uint64_t));
        }
        if (typ->id() == Type::NA) {
            return KeyColumnMetadata(true, 0, true);
        }
        // Caller attempted to create a KeyColumnArray from an invalid type
        return turbo::failed_precondition_error("Unsupported column data type ", typ->name(),
                                                " used with KeyColumnMetadata");
    }

    turbo::Result<KeyColumnArray> ColumnArrayFromArrayData(
            const std::shared_ptr<ArrayData> &array_data, int64_t start_row, int64_t num_rows) {
        TURBO_MOVE_OR_RAISE(KeyColumnMetadata metadata,
                            ColumnMetadataFromDataType(array_data->type));
        return ColumnArrayFromArrayDataAndMetadata(array_data, metadata, start_row, num_rows);
    }

    KeyColumnArray ColumnArrayFromArrayDataAndMetadata(
            const std::shared_ptr<ArrayData> &array_data, const KeyColumnMetadata &metadata,
            int64_t start_row, int64_t num_rows) {
        KeyColumnArray column_array = KeyColumnArray(
                metadata, array_data->offset + start_row + num_rows,
                array_data->buffers[0] != nullptr ? array_data->buffers[0]->data() : nullptr,
                array_data->buffers[1]->data(),
                (array_data->buffers.size() > 2 && array_data->buffers[2] != nullptr)
                ? array_data->buffers[2]->data()
                : nullptr);
        return column_array.slice(array_data->offset + start_row, num_rows);
    }

    turbo::Status ColumnMetadatasFromExecBatch(const ExecBatch &batch,
                                               std::vector<KeyColumnMetadata> *column_metadatas) {
        int num_columns = static_cast<int>(batch.values.size());
        column_metadatas->resize(num_columns);
        for (int i = 0; i < num_columns; ++i) {
            const Datum &data = batch.values[i];
                    DKCHECK(data.is_array());
            const std::shared_ptr<ArrayData> &array_data = data.array();
            TURBO_MOVE_OR_RAISE((*column_metadatas)[i],
                                ColumnMetadataFromDataType(array_data->type));
        }
        return turbo::OkStatus();
    }

    turbo::Status ColumnArraysFromExecBatch(const ExecBatch &batch, int64_t start_row,
                                            int64_t num_rows,
                                            std::vector<KeyColumnArray> *column_arrays) {
        int num_columns = static_cast<int>(batch.values.size());
        column_arrays->resize(num_columns);
        for (int i = 0; i < num_columns; ++i) {
            const Datum &data = batch.values[i];
                    DKCHECK(data.is_array());
            const std::shared_ptr<ArrayData> &array_data = data.array();
            TURBO_MOVE_OR_RAISE((*column_arrays)[i],
                                ColumnArrayFromArrayData(array_data, start_row, num_rows));
        }
        return turbo::OkStatus();
    }

    turbo::Status ColumnArraysFromExecBatch(const ExecBatch &batch,
                                            std::vector<KeyColumnArray> *column_arrays) {
        return ColumnArraysFromExecBatch(batch, 0, static_cast<int>(batch.length),
                                         column_arrays);
    }

    void ResizableArrayData::init(const std::shared_ptr<DataType> &data_type,
                                  MemoryPool *pool, int log_num_rows_min) {
#ifndef NDEBUG
        if (num_rows_allocated_ > 0) {
                    DKCHECK(data_type_ != nullptr);
            KeyColumnMetadata metadata_before =
                    ColumnMetadataFromDataType(data_type_).value_or_die();
            KeyColumnMetadata metadata_after = ColumnMetadataFromDataType(data_type).value_or_die();
                    DKCHECK(metadata_before.is_fixed_length == metadata_after.is_fixed_length &&
                                  metadata_before.fixed_length == metadata_after.fixed_length);
        }
#endif
        Clear(/*release_buffers=*/false);
        log_num_rows_min_ = log_num_rows_min;
        data_type_ = data_type;
        pool_ = pool;
    }

    void ResizableArrayData::Clear(bool release_buffers) {
        num_rows_ = 0;
        if (release_buffers) {
            buffers_[kValidityBuffer].reset();
            buffers_[kFixedLengthBuffer].reset();
            buffers_[kVariableLengthBuffer].reset();
            num_rows_allocated_ = 0;
            var_len_buf_size_ = 0;
        }
    }

    turbo::Status ResizableArrayData::ResizeFixedLengthBuffers(int num_rows_new) {
                DKCHECK(num_rows_new >= 0);
        if (num_rows_new <= num_rows_allocated_) {
            num_rows_ = num_rows_new;
            return turbo::OkStatus();
        }

        int num_rows_allocated_new = 1 << log_num_rows_min_;
        while (num_rows_allocated_new < num_rows_new) {
            num_rows_allocated_new *= 2;
        }

        KeyColumnMetadata column_metadata = ColumnMetadataFromDataType(data_type_).value_or_die();

        if (buffers_[kFixedLengthBuffer] == nullptr) {
                    DKCHECK(buffers_[kValidityBuffer] == nullptr &&
                                  buffers_[kVariableLengthBuffer] == nullptr);

            TURBO_MOVE_OR_RAISE(
                    buffers_[kValidityBuffer],
                    allocate_resizable_buffer(
                            bit_util::BytesForBits(num_rows_allocated_new) + kNumPaddingBytes, pool_));
            memset(mutable_data(kValidityBuffer), 0,
                   bit_util::BytesForBits(num_rows_allocated_new) + kNumPaddingBytes);
            if (column_metadata.is_fixed_length) {
                if (column_metadata.fixed_length == 0) {
                    TURBO_MOVE_OR_RAISE(
                            buffers_[kFixedLengthBuffer],
                            allocate_resizable_buffer(
                                    bit_util::BytesForBits(num_rows_allocated_new) + kNumPaddingBytes,
                                    pool_));
                    memset(mutable_data(kFixedLengthBuffer), 0,
                           bit_util::BytesForBits(num_rows_allocated_new) + kNumPaddingBytes);
                } else {
                    TURBO_MOVE_OR_RAISE(
                            buffers_[kFixedLengthBuffer],
                            allocate_resizable_buffer(
                                    num_rows_allocated_new * column_metadata.fixed_length + kNumPaddingBytes,
                                    pool_));
                }
            } else {
                TURBO_MOVE_OR_RAISE(
                        buffers_[kFixedLengthBuffer],
                        allocate_resizable_buffer(
                                (num_rows_allocated_new + 1) * sizeof(uint32_t) + kNumPaddingBytes, pool_));
            }

            TURBO_MOVE_OR_RAISE(
                    buffers_[kVariableLengthBuffer],
                    allocate_resizable_buffer(sizeof(uint64_t) + kNumPaddingBytes, pool_));

            var_len_buf_size_ = sizeof(uint64_t);
        } else {
                    DKCHECK(buffers_[kValidityBuffer] != nullptr &&
                                  buffers_[kVariableLengthBuffer] != nullptr);

            int64_t bytes_for_bits_before =
                    bit_util::BytesForBits(num_rows_allocated_) + kNumPaddingBytes;
            int64_t bytes_for_bits_after =
                    bit_util::BytesForBits(num_rows_allocated_new) + kNumPaddingBytes;

            TURBO_RETURN_NOT_OK(buffers_[kValidityBuffer]->resize(
                    bit_util::BytesForBits(num_rows_allocated_new) + kNumPaddingBytes));
            memset(mutable_data(kValidityBuffer) + bytes_for_bits_before, 0,
                   bytes_for_bits_after - bytes_for_bits_before);

            if (column_metadata.is_fixed_length) {
                if (column_metadata.fixed_length == 0) {
                    TURBO_RETURN_NOT_OK(buffers_[kFixedLengthBuffer]->resize(
                            bit_util::BytesForBits(num_rows_allocated_new) + kNumPaddingBytes));
                    memset(mutable_data(kFixedLengthBuffer) + bytes_for_bits_before, 0,
                           bytes_for_bits_after - bytes_for_bits_before);
                } else {
                    TURBO_RETURN_NOT_OK(buffers_[kFixedLengthBuffer]->resize(
                            num_rows_allocated_new * column_metadata.fixed_length + kNumPaddingBytes));
                }
            } else {
                TURBO_RETURN_NOT_OK(buffers_[kFixedLengthBuffer]->resize(
                        (num_rows_allocated_new + 1) * sizeof(uint32_t) + kNumPaddingBytes));
            }
        }

        num_rows_allocated_ = num_rows_allocated_new;
        num_rows_ = num_rows_new;

        return turbo::OkStatus();
    }

    turbo::Status ResizableArrayData::ResizeVaryingLengthBuffer() {
        KeyColumnMetadata column_metadata;
        column_metadata = ColumnMetadataFromDataType(data_type_).value_or_die();

        if (!column_metadata.is_fixed_length) {
            int64_t min_new_size = buffers_[kFixedLengthBuffer]->data_as<int32_t>()[num_rows_];
                    DKCHECK(var_len_buf_size_ > 0);
            if (var_len_buf_size_ < min_new_size) {
                int64_t new_size = var_len_buf_size_;
                while (new_size < min_new_size) {
                    new_size *= 2;
                }
                TURBO_RETURN_NOT_OK(buffers_[kVariableLengthBuffer]->resize(new_size + kNumPaddingBytes));
                var_len_buf_size_ = new_size;
            }
        }

        return turbo::OkStatus();
    }

    KeyColumnArray ResizableArrayData::column_array() const {
        KeyColumnMetadata column_metadata;
        column_metadata = ColumnMetadataFromDataType(data_type_).value_or_die();
        return KeyColumnArray(column_metadata, num_rows_,
                              buffers_[kValidityBuffer]->mutable_data(),
                              buffers_[kFixedLengthBuffer]->mutable_data(),
                              buffers_[kVariableLengthBuffer]->mutable_data());
    }

    std::shared_ptr<ArrayData> ResizableArrayData::array_data() const {
        KeyColumnMetadata column_metadata;
        column_metadata = ColumnMetadataFromDataType(data_type_).value_or_die();

        auto valid_count = nebula::internal::CountSetBits(
                buffers_[kValidityBuffer]->data(), /*offset=*/0, static_cast<int64_t>(num_rows_));
        int null_count = static_cast<int>(num_rows_) - static_cast<int>(valid_count);

        if (column_metadata.is_fixed_length) {
            return ArrayData::create(data_type_, num_rows_,
                                   {buffers_[kValidityBuffer], buffers_[kFixedLengthBuffer]},
                                   null_count);
        } else {
            return ArrayData::create(data_type_, num_rows_,
                                   {buffers_[kValidityBuffer], buffers_[kFixedLengthBuffer],
                                    buffers_[kVariableLengthBuffer]},
                                   null_count);
        }
    }

    int ExecBatchBuilder::NumRowsToSkip(const std::shared_ptr<ArrayData> &column,
                                        int num_rows, const uint16_t *row_ids,
                                        int num_tail_bytes_to_skip) {
#ifndef NDEBUG
        // Ids must be in non-decreasing order
        //
        for (int i = 1; i < num_rows; ++i) {
                    DKCHECK(row_ids[i] >= row_ids[i - 1]);
        }
#endif

        KeyColumnMetadata column_metadata =
                ColumnMetadataFromDataType(column->type).value_or_die();
                DKCHECK(!column_metadata.is_fixed_length || column_metadata.fixed_length > 0);

        int num_rows_left = num_rows;
        int num_bytes_skipped = 0;
        while (num_rows_left > 0 && num_bytes_skipped < num_tail_bytes_to_skip) {
            --num_rows_left;
            int row_id_removed = row_ids[num_rows_left];
            if (column_metadata.is_fixed_length) {
                num_bytes_skipped += column_metadata.fixed_length;
            } else {
                const int32_t *offsets = column->get_values<int32_t>(1);
                num_bytes_skipped += offsets[row_id_removed + 1] - offsets[row_id_removed];
            }
            // Skip consecutive rows with the same id
            while (num_rows_left > 0 && row_id_removed == row_ids[num_rows_left - 1]) {
                --num_rows_left;
            }
        }

        return num_rows - num_rows_left;
    }

    template<bool OUTPUT_BYTE_ALIGNED>
    void ExecBatchBuilder::CollectBitsImp(const uint8_t *input_bits,
                                          int64_t input_bits_offset, uint8_t *output_bits,
                                          int64_t output_bits_offset, int num_rows,
                                          const uint16_t *row_ids) {
        if (!OUTPUT_BYTE_ALIGNED) {
                    DKCHECK(output_bits_offset % 8 > 0);
            output_bits[output_bits_offset / 8] &=
                    static_cast<uint8_t>((1 << (output_bits_offset % 8)) - 1);
        } else {
                    DKCHECK(output_bits_offset % 8 == 0);
        }
        constexpr int unroll = 8;
        for (int i = 0; i < num_rows / unroll; ++i) {
            const uint16_t *row_ids_base = row_ids + unroll * i;
            uint8_t result;
            result = bit_util::get_bit(input_bits, input_bits_offset + row_ids_base[0]) ? 1 : 0;
            result |= bit_util::get_bit(input_bits, input_bits_offset + row_ids_base[1]) ? 2 : 0;
            result |= bit_util::get_bit(input_bits, input_bits_offset + row_ids_base[2]) ? 4 : 0;
            result |= bit_util::get_bit(input_bits, input_bits_offset + row_ids_base[3]) ? 8 : 0;
            result |= bit_util::get_bit(input_bits, input_bits_offset + row_ids_base[4]) ? 16 : 0;
            result |= bit_util::get_bit(input_bits, input_bits_offset + row_ids_base[5]) ? 32 : 0;
            result |= bit_util::get_bit(input_bits, input_bits_offset + row_ids_base[6]) ? 64 : 0;
            result |= bit_util::get_bit(input_bits, input_bits_offset + row_ids_base[7]) ? 128 : 0;
            if (OUTPUT_BYTE_ALIGNED) {
                output_bits[output_bits_offset / 8 + i] = result;
            } else {
                output_bits[output_bits_offset / 8 + i] |=
                        static_cast<uint8_t>(result << (output_bits_offset % 8));
                output_bits[output_bits_offset / 8 + i + 1] =
                        static_cast<uint8_t>(result >> (8 - (output_bits_offset % 8)));
            }
        }
        if (num_rows % unroll > 0) {
            for (int i = num_rows - (num_rows % unroll); i < num_rows; ++i) {
                bit_util::SetBitTo(output_bits, output_bits_offset + i,
                                   bit_util::get_bit(input_bits, input_bits_offset + row_ids[i]));
            }
        }
    }

    void ExecBatchBuilder::CollectBits(const uint8_t *input_bits, int64_t input_bits_offset,
                                       uint8_t *output_bits, int64_t output_bits_offset,
                                       int num_rows, const uint16_t *row_ids) {
        if (output_bits_offset % 8 > 0) {
            CollectBitsImp<false>(input_bits, input_bits_offset, output_bits, output_bits_offset,
                                  num_rows, row_ids);
        } else {
            CollectBitsImp<true>(input_bits, input_bits_offset, output_bits, output_bits_offset,
                                 num_rows, row_ids);
        }
    }

    template<class PROCESS_VALUE_FN>
    void ExecBatchBuilder::Visit(const std::shared_ptr<ArrayData> &column, int num_rows,
                                 const uint16_t *row_ids, PROCESS_VALUE_FN process_value_fn) {
        KeyColumnMetadata metadata = ColumnMetadataFromDataType(column->type).value_or_die();

        if (!metadata.is_fixed_length) {
            const uint8_t *ptr_base = column->buffers[2]->data();
            const int32_t *offsets = column->get_values<int32_t>(1);
            for (int i = 0; i < num_rows; ++i) {
                uint16_t row_id = row_ids[i];
                const uint8_t *field_ptr = ptr_base + offsets[row_id];
                int32_t field_length = offsets[row_id + 1] - offsets[row_id];
                process_value_fn(i, field_ptr, field_length);
            }
        } else {
                    DKCHECK(metadata.fixed_length > 0);
            for (int i = 0; i < num_rows; ++i) {
                uint16_t row_id = row_ids[i];
                const uint8_t *field_ptr =
                        column->buffers[1]->data() +
                        (column->offset + row_id) * static_cast<int64_t>(metadata.fixed_length);
                process_value_fn(i, field_ptr, static_cast<int32_t>(metadata.fixed_length));
            }
        }
    }

    turbo::Status ExecBatchBuilder::AppendSelected(const std::shared_ptr<ArrayData> &source,
                                                   ResizableArrayData *target,
                                                   int num_rows_to_append, const uint16_t *row_ids,
                                                   MemoryPool *pool) {
        int num_rows_before = target->num_rows();
                DKCHECK(num_rows_before >= 0);
        int num_rows_after = num_rows_before + num_rows_to_append;
        if (target->num_rows() == 0) {
            target->init(source->type, pool, kLogNumRows);
        }
        TURBO_RETURN_NOT_OK(target->ResizeFixedLengthBuffers(num_rows_after));

        KeyColumnMetadata column_metadata =
                ColumnMetadataFromDataType(source->type).value_or_die();

        if (column_metadata.is_fixed_length) {
            // Fixed length column
            //
            uint32_t fixed_length = column_metadata.fixed_length;
            switch (fixed_length) {
                case 0:
                    CollectBits(source->buffers[1]->data(), source->offset, target->mutable_data(1),
                                num_rows_before, num_rows_to_append, row_ids);
                    break;
                case 1:
                    Visit(source, num_rows_to_append, row_ids,
                          [&](int i, const uint8_t *ptr, int32_t num_bytes) {
                              target->mutable_data(1)[num_rows_before + i] = *ptr;
                          });
                    break;
                case 2:
                    Visit(
                            source, num_rows_to_append, row_ids,
                            [&](int i, const uint8_t *ptr, int32_t num_bytes) {
                                reinterpret_cast<uint16_t *>(target->mutable_data(1))[num_rows_before + i] =
                                        *reinterpret_cast<const uint16_t *>(ptr);
                            });
                    break;
                case 4:
                    Visit(
                            source, num_rows_to_append, row_ids,
                            [&](int i, const uint8_t *ptr, int32_t num_bytes) {
                                reinterpret_cast<uint32_t *>(target->mutable_data(1))[num_rows_before + i] =
                                        *reinterpret_cast<const uint32_t *>(ptr);
                            });
                    break;
                case 8:
                    Visit(
                            source, num_rows_to_append, row_ids,
                            [&](int i, const uint8_t *ptr, int32_t num_bytes) {
                                reinterpret_cast<uint64_t *>(target->mutable_data(1))[num_rows_before + i] =
                                        *reinterpret_cast<const uint64_t *>(ptr);
                            });
                    break;
                default: {
                    int num_rows_to_process =
                            num_rows_to_append -
                            NumRowsToSkip(source, num_rows_to_append, row_ids, sizeof(uint64_t));
                    Visit(source, num_rows_to_process, row_ids,
                          [&](int i, const uint8_t *ptr, int32_t num_bytes) {
                              uint64_t *dst = reinterpret_cast<uint64_t *>(
                                      target->mutable_data(1) +
                                      static_cast<int64_t>(num_bytes) * (num_rows_before + i));
                              const uint64_t *src = reinterpret_cast<const uint64_t *>(ptr);
                              for (uint32_t word_id = 0;
                                   word_id < bit_util::CeilDiv(num_bytes, sizeof(uint64_t));
                                   ++word_id) {
                                  turbo::safe_store<uint64_t>(dst + word_id, turbo::safe_load(src + word_id));
                              }
                          });
                    if (num_rows_to_append > num_rows_to_process) {
                        Visit(source, num_rows_to_append - num_rows_to_process,
                              row_ids + num_rows_to_process,
                              [&](int i, const uint8_t *ptr, int32_t num_bytes) {
                                  uint64_t *dst = reinterpret_cast<uint64_t *>(
                                          target->mutable_data(1) +
                                          static_cast<int64_t>(num_bytes) *
                                          (num_rows_before + num_rows_to_process + i));
                                  const uint64_t *src = reinterpret_cast<const uint64_t *>(ptr);
                                  memcpy(dst, src, num_bytes);
                              });
                    }
                }
            }
        } else {
            // Varying length column
            //

            // Step 1: calculate target offsets
            //
            int32_t *offsets = reinterpret_cast<int32_t *>(target->mutable_data(1));
            int32_t sum = num_rows_before == 0 ? 0 : offsets[num_rows_before];
            Visit(source, num_rows_to_append, row_ids,
                  [&](int i, const uint8_t *ptr, int32_t num_bytes) {
                      offsets[num_rows_before + i] = num_bytes;
                  });
            for (int i = 0; i < num_rows_to_append; ++i) {
                int32_t length = offsets[num_rows_before + i];
                offsets[num_rows_before + i] = sum;
                int32_t new_sum_maybe_overflow = 0;
                if (TURBO_UNLIKELY(
                        nebula::internal::AddWithOverflow(sum, length, &new_sum_maybe_overflow))) {
                    return turbo::invalid_argument_error("Overflow detected in ExecBatchBuilder when appending ",
                                                         num_rows_before + i + 1, "-th element of length ", length,
                                                         " bytes to current length ", sum, " bytes");
                }
                sum = new_sum_maybe_overflow;
            }
            offsets[num_rows_before + num_rows_to_append] = sum;

            // Step 2: resize output buffers
            //
            TURBO_RETURN_NOT_OK(target->ResizeVaryingLengthBuffer());

            // Step 3: copy varying-length data
            //
            int num_rows_to_process =
                    num_rows_to_append -
                    NumRowsToSkip(source, num_rows_to_append, row_ids, sizeof(uint64_t));
            Visit(source, num_rows_to_process, row_ids,
                  [&](int i, const uint8_t *ptr, int32_t num_bytes) {
                      uint64_t *dst = reinterpret_cast<uint64_t *>(target->mutable_data(2) +
                                                                   offsets[num_rows_before + i]);
                      const uint64_t *src = reinterpret_cast<const uint64_t *>(ptr);
                      for (uint32_t word_id = 0;
                           word_id < bit_util::CeilDiv(num_bytes, sizeof(uint64_t)); ++word_id) {
                          turbo::safe_store<uint64_t>(dst + word_id, turbo::safe_load(src + word_id));
                      }
                  });
            Visit(source, num_rows_to_append - num_rows_to_process, row_ids + num_rows_to_process,
                  [&](int i, const uint8_t *ptr, int32_t num_bytes) {
                      uint64_t *dst = reinterpret_cast<uint64_t *>(
                              target->mutable_data(2) +
                              offsets[num_rows_before + num_rows_to_process + i]);
                      const uint64_t *src = reinterpret_cast<const uint64_t *>(ptr);
                      memcpy(dst, src, num_bytes);
                  });
        }

        // process nulls
        //
        if (source->buffers[0] == nullptr) {
            uint8_t *dst = target->mutable_data(0);
            dst[num_rows_before / 8] |= static_cast<uint8_t>(~0ULL << (num_rows_before & 7));
            for (int i = num_rows_before / 8 + 1;
                 i < bit_util::BytesForBits(num_rows_before + num_rows_to_append); ++i) {
                dst[i] = 0xff;
            }
        } else {
            CollectBits(source->buffers[0]->data(), source->offset, target->mutable_data(0),
                        num_rows_before, num_rows_to_append, row_ids);
        }

        return turbo::OkStatus();
    }

    turbo::Status ExecBatchBuilder::append_nulls(const std::shared_ptr<DataType> &type,
                                                ResizableArrayData &target, int num_rows_to_append,
                                                MemoryPool *pool) {
        int num_rows_before = target.num_rows();
        int num_rows_after = num_rows_before + num_rows_to_append;
        if (target.num_rows() == 0) {
            target.init(type, pool, kLogNumRows);
        }
        TURBO_RETURN_NOT_OK(target.ResizeFixedLengthBuffers(num_rows_after));

        KeyColumnMetadata column_metadata = ColumnMetadataFromDataType(type).value_or_die();

        // process fixed length buffer
        //
        if (column_metadata.is_fixed_length) {
            uint8_t *dst = target.mutable_data(1);
            if (column_metadata.fixed_length == 0) {
                dst[num_rows_before / 8] &= static_cast<uint8_t>((1 << (num_rows_before % 8)) - 1);
                int64_t offset_begin = num_rows_before / 8 + 1;
                int64_t offset_end = bit_util::BytesForBits(num_rows_after);
                if (offset_end > offset_begin) {
                    memset(dst + offset_begin, 0, offset_end - offset_begin);
                }
            } else {
                memset(dst + num_rows_before * static_cast<int64_t>(column_metadata.fixed_length),
                       0, static_cast<int64_t>(column_metadata.fixed_length) * num_rows_to_append);
            }
        } else {
            uint32_t *dst = reinterpret_cast<uint32_t *>(target.mutable_data(1));
            uint32_t sum = num_rows_before == 0 ? 0 : dst[num_rows_before];
            for (int64_t i = num_rows_before; i <= num_rows_after; ++i) {
                dst[i] = sum;
            }
        }

        // process nulls
        //
        uint8_t *dst = target.mutable_data(0);
        dst[num_rows_before / 8] &= static_cast<uint8_t>((1 << (num_rows_before % 8)) - 1);
        int64_t offset_begin = num_rows_before / 8 + 1;
        int64_t offset_end = bit_util::BytesForBits(num_rows_after);
        if (offset_end > offset_begin) {
            memset(dst + offset_begin, 0, offset_end - offset_begin);
        }

        return turbo::OkStatus();
    }

    turbo::Status ExecBatchBuilder::AppendSelected(MemoryPool *pool, const ExecBatch &batch,
                                                   int num_rows_to_append, const uint16_t *row_ids,
                                                   int num_cols, const int *col_ids) {
        if (num_rows_to_append == 0) {
            return turbo::OkStatus();
        }

        if (num_rows() + num_rows_to_append > num_rows_max()) {
            return turbo::resource_exhausted_error("ExecBatch builder exceeded limit of accumulated rows");
        }

        // If this is the first time we append rows, then initialize output buffers.
        //
        if (values_.empty()) {
            values_.resize(num_cols);
            for (int i = 0; i < num_cols; ++i) {
                const Datum &data = batch.values[col_ids ? col_ids[i] : i];
                        DKCHECK(data.is_array());
                const std::shared_ptr<ArrayData> &array_data = data.array();
                values_[i].init(array_data->type, pool, kLogNumRows);
            }
        }

        for (size_t i = 0; i < values_.size(); ++i) {
            const Datum &data = batch.values[col_ids ? col_ids[i] : i];
                    DKCHECK(data.is_array());
            const std::shared_ptr<ArrayData> &array_data = data.array();
            TURBO_RETURN_NOT_OK(
                    AppendSelected(array_data, &values_[i], num_rows_to_append, row_ids, pool));
        }

        return turbo::OkStatus();
    }

    turbo::Status ExecBatchBuilder::append_nulls(MemoryPool *pool,
                                                const std::vector<std::shared_ptr<DataType>> &types,
                                                int num_rows_to_append) {
        if (num_rows_to_append == 0) {
            return turbo::OkStatus();
        }

        if (num_rows() + num_rows_to_append > num_rows_max()) {
            return turbo::resource_exhausted_error("ExecBatch builder exceeded limit of accumulated rows.");
        }

        // If this is the first time we append rows, then initialize output buffers.
        //
        if (values_.empty()) {
            values_.resize(types.size());
            for (size_t i = 0; i < types.size(); ++i) {
                values_[i].init(types[i], pool, kLogNumRows);
            }
        }

        for (size_t i = 0; i < values_.size(); ++i) {
            TURBO_RETURN_NOT_OK(append_nulls(types[i], values_[i], num_rows_to_append, pool));
        }

        return turbo::OkStatus();
    }

    ExecBatch ExecBatchBuilder::flush() {
                DKCHECK(num_rows() > 0);
        ExecBatch out({}, num_rows());
        out.values.resize(values_.size());
        for (size_t i = 0; i < values_.size(); ++i) {
            out.values[i] = values_[i].array_data();
            values_[i].Clear(true);
        }
        return out;
    }

}  // namespace nebula::compute
